from fastmcp import FastMCP
from typing import Union
from Tools.MySQLClient import MySQLClient, MySQLManager
from dotenv import load_dotenv
import os


load_dotenv()


mcp = FastMCP(
    name="MySQL_MCP",
    instructions="""
    MySQL MCP 是一个用于管理 MySQL 数据库的 MCP 工具。
    """,
    host=os.getenv("HOST"),
    port=int(os.getenv("PORT")),
)



@mcp.tool(name="get_databases",description="获取数据库列表")
def get_databases(mysql_manager: dict) -> list:
    """
    获取MySQL服务器上的所有数据库列表
    
    Args:
        mysql_manager: 包含MySQL连接信息的字典，包括host, port, user, password等
                       示例：{"host": "localhost", "port": 3306, "user": "root", "password": "123456"}
        
    Returns:
        list: 数据库名称列表
    """
    try:
        # 从字典中提取连接参数
        host = mysql_manager.get("host", "localhost")
        port = mysql_manager.get("port", 3306)
        user = mysql_manager.get("user")
        password = mysql_manager.get("password")
        database = mysql_manager.get("database")   # 可以为None
        
        # 创建MySQLManager实例
        manager = MySQLManager(host=host, port=port, user=user, password=password, database=database)
        
        # 连接到MySQL
        if not manager.connect():
            return {"error": "无法连接到MySQL服务器"}
        
        # 使用MySQLClient获取数据库列表
        client = MySQLClient()
        databases = client.get_databases(manager)
        
        # 关闭连接
        manager.connection.close()
        
        return databases
    except Exception as e:
        return {"error": f"获取数据库列表失败: {str(e)}"}

@mcp.tool(name="get_tables", description="获取数据库中的表列表")
def get_tables(mysql_manager: dict, database: str = None) -> dict:
    """
    获取数据库中的表列表
    
    Args:
        mysql_manager: 包含MySQL连接信息的字典，示例：{"host": "localhost", "port": 3306, "user": "root", "password": "123456"}
        database: 数据库名称，如果为None则使用连接时设置的默认数据库
        
    Returns:
        dict: 表列表或错误信息
    """
    try:
        # 从字典中提取连接参数
        host = mysql_manager.get("host", "localhost")
        port = mysql_manager.get("port", 3306)
        user = mysql_manager.get("user")
        password = mysql_manager.get("password")
        default_database = mysql_manager.get("database")  # 可以为None
        
        # 创建MySQLManager实例
        manager = MySQLManager(host=host, port=port, user=user, password=password, database=default_database)
        
        # 连接到MySQL
        if not manager.connect():
            return {"error": "无法连接到MySQL服务器"}
        
        # 使用MySQLClient获取表列表
        client = MySQLClient()
        tables = client.get_tables(manager, database)
        
        # 关闭连接
        manager.connection.close()
        
        # 检查是否返回了错误
        if isinstance(tables, dict) and "error" in tables:
            return tables
        
        return {"tables": tables}
    except Exception as e:
        return {"error": f"获取表列表失败: {str(e)}"}

@mcp.tool(name="get_table_structure", description="获取表结构信息")
def get_table_structure(mysql_manager: dict, table: Union[str, list] = None, database: str = None) -> dict:
    """
    获取表结构信息
    
    Args:
        mysql_manager: 包含MySQL连接信息的字典，示例：{"host": "localhost", "port": 3306, "user": "root", "password": "123456"}
        table: 表名（字符串）或表名列表，如果为None则获取所有表的结构，否则获取指定表的结构，示例："table1" 或 ["table1", "table2"]
        database: 数据库名称，如果为None则使用连接时设置的默认数据库
        
    Returns:
        dict: 表结构信息或错误信息
    """
    try:
        # 从字典中提取连接参数
        host = mysql_manager.get("host", "localhost")
        port = mysql_manager.get("port", 3306)
        user = mysql_manager.get("user")
        password = mysql_manager.get("password")
        default_database = mysql_manager.get("database")  # 可以为None
        
        # 创建MySQLManager实例
        manager = MySQLManager(host=host, port=port, user=user, password=password, database=default_database)
        
        # 连接到MySQL
        if not manager.connect():
            return {"error": "无法连接到MySQL服务器"}
        
        # 使用MySQLClient获取表结构
        client = MySQLClient()
        
        # 处理table参数，如果它是JSON字符串，则解析为列表
        import json
        if table is not None and isinstance(table, str):
            try:
                # 尝试解析JSON字符串
                parsed_table = json.loads(table)
                if isinstance(parsed_table, list):
                    table = parsed_table
            except (json.JSONDecodeError, TypeError):
                # 如果不是有效的JSON，则保持原样（作为单个表名）
                pass
        
        structure = client.get_table_structure(manager, table, database)
        
        # 关闭连接
        manager.connection.close()
        
        # 检查是否返回了错误
        if isinstance(structure, dict) and "error" in structure:
            return structure
        
        return {"structure": structure}
    except Exception as e:
        return {"error": f"获取表结构失败: {str(e)}"}

@mcp.tool(name="execute_query", description="执行SQL查询语句")
def execute_query(mysql_manager: dict, query: str, database: str = None) -> dict:
    """
    执行SQL查询语句
    
    Args:
        mysql_manager: 包含MySQL连接信息的字典，示例：{"host": "localhost", "port": 3306, "user": "root", "password": "123456"}
        query: SQL查询语句
        database: 数据库名称，如果为None则使用连接时设置的默认数据库
        
    Returns:
        dict: 查询结果或错误信息
    """
    try:
        # 从字典中提取连接参数
        host = mysql_manager.get("host", "localhost")
        port = mysql_manager.get("port", 3306)
        user = mysql_manager.get("user")
        password = mysql_manager.get("password")
        default_database = mysql_manager.get("database")  # 可以为None
        
        # 创建MySQLManager实例
        manager = MySQLManager(host=host, port=port, user=user, password=password, database=default_database)
        
        # 连接到MySQL
        if not manager.connect():
            return {"error": "无法连接到MySQL服务器"}
        
        # 使用MySQLClient执行查询
        client = MySQLClient()
        result = client.execute_query(manager, query, database)
        
        # 关闭连接
        manager.connection.close()
        
        # 检查是否返回了错误
        if isinstance(result, dict) and "error" in result:
            return result
        
        return {"result": result}
    except Exception as e:
        return {"error": f"执行查询失败: {str(e)}"}

@mcp.tool(name="execute_non_query", description="执行非查询SQL语句（INSERT、UPDATE、DELETE等），支持批量操作")
def execute_non_query(mysql_manager: dict, query: str, database: str = None, params: list = None) -> dict:
    """
    执行非查询SQL语句（INSERT、UPDATE、DELETE等），支持批量操作
    
    Args:
        mysql_manager: 包含MySQL连接信息的字典，示例：{"host": "localhost", "port": 3306, "user": "root", "password": "123456"}
        query: SQL非查询语句
        database: 数据库名称，如果为None则使用连接时设置的默认数据库
        params: 参数化查询的参数列表
               - 单个操作时为元组或列表，如 ["John", "john@example.com"]
               - 批量操作时为元组列表或列表的列表，如 [["John", "john@example.com"], ["Jane", "jane@example.com"]]
        
    Returns:
        dict: 执行结果或错误信息，包含affected_rows（影响行数）和batch_count（批量操作时的批次数）
    """
    try:
        # 从字典中提取连接参数
        host = mysql_manager.get("host", "localhost")
        port = mysql_manager.get("port", 3306)
        user = mysql_manager.get("user")
        password = mysql_manager.get("password")
        default_database = mysql_manager.get("database")  # 可以为None
        
        # 创建MySQLManager实例
        manager = MySQLManager(host=host, port=port, user=user, password=password, database=default_database)
        
        # 连接到MySQL
        if not manager.connect():
            return {"error": "无法连接到MySQL服务器"}
        
        # 使用MySQLClient执行非查询
        client = MySQLClient()
        
        # 自动检测是否为批量操作
        is_batch = False
        if params and len(params) > 0:
            # 如果params的第一个元素是列表或元组，且params本身不是字符串，则认为是批量操作
            if isinstance(params[0], (list, tuple)) and not isinstance(params, str):
                is_batch = True
        
        if is_batch:
            # 批量操作模式
            total_affected_rows = 0
            batch_count = len(params)
            
            # 执行批量操作
            for param_tuple in params:
                # 确保每个参数组是元组类型
                if not isinstance(param_tuple, tuple):
                    param_tuple = tuple(param_tuple)
                result = client.execute_non_query(manager, query, database, param_tuple)
                
                # 检查是否返回了错误
                if isinstance(result, dict) and "error" in result:
                    return result
                
                total_affected_rows += result
            
            # 关闭连接
            manager.connection.close()
            
            return {
                "affected_rows": total_affected_rows,
                "batch_count": batch_count,
                "message": f"成功执行批量操作，共{batch_count}批，影响{total_affected_rows}行"
            }
        else:
            # 单个操作模式
            # 将params转换为tuple类型，因为MySQLClient.execute_non_query需要tuple类型
            params_tuple = tuple(params) if params else None
            result = client.execute_non_query(manager, query, database, params_tuple)
            
            # 检查是否返回了错误
            if isinstance(result, dict) and "error" in result:
                return result
            
            manager.connection.close()
            return {"affected_rows": result}
        
        # 关闭连接
        
    except Exception as e:
        manager.connection.close()
        return {"error": f"执行非查询失败: {str(e)}"}

if __name__ == "__main__":
    mcp.run("sse")
